跳转至

基于自适应注意力匹配的推理过程思维感知 KV 缓存压缩

文章背景与核心概要

推理型语言模型在生成复杂的思维链（CoT）序列时，会产生大量的键值（KV）缓存，这在解码阶段会造成严重的内存瓶颈。传统的压缩方法通常将推理轨迹视为平坦的 Token 序列并应用统一的压缩策略，忽略了推理步骤在重要性上存在显著差异的层次化结构。

为了解决这一问题，作者提出了思维感知注意力匹配（TAM）方法，该方法通过三个核心机制利用结构化推理动态：(1) 思维分割，将轨迹分解为不同的推理块；(2) 自适应预算分配，根据每个段的大小和相对重要性分配合适的压缩预算；(3) 关键 Token 保护，以保留高注意力的推理锚点。理论证明表明，该分配规则在凸误差模型下保持最优，且顺序压缩下的累积误差具有严格的边界。在 AIME 2024 和 MATH-500 基准测试中基于 Qwen3-4B 的实证评估表明，在相同的内存占用下，TAM 能够提高任务准确率；此外，周期性压缩将峰值内存消耗限制在 3.1–3.2 GB（减少了 65%），同时保持了具有竞争力的准确率。

Executive Summary / 执行摘要

Reasoning language models generate extensive chain-of-thought (CoT) sequences, leading to a linear growth of the key-value (KV) cache that acts as a significant memory bottleneck during decoding. Conventional compaction methods typically treat reasoning trajectories as flat token sequences and apply uniform compression, failing to account for the hierarchical structure where reasoning steps vary significantly in importance.

To overcome this, the authors introduce Thought-Aware Attention Matching (TAM), a novel approach leveraging structural reasoning dynamics through three core mechanisms: 1. Thought segmentation to decompose trajectories into distinct reasoning blocks. 2. Adaptive budget allocation to distribute compression budgets according to each segment's size and relative importance. 3. Pivotal token protection to preserve critical high-attention reasoning anchors.

Theoretical proofs demonstrate that the allocation rule remains optimal under a convex error model, and cumulative errors under sequential compaction are strictly bounded. Empirical evaluations on AIME 2024 and MATH-500 utilizing Qwen3-4B show that TAM improves task accuracy compared to uniform compaction at equivalent memory footprints. Furthermore, periodic compaction bounds peak memory consumption to 3.1–3.2 GB (a 65% reduction) while sustaining competitive accuracy.

推理语言模型会生成大量的思维链（CoT）序列，导致键值（KV）缓存呈线性增长，这成为解码过程中的重大内存瓶颈。常规的压缩方法通常将推理轨迹视为平坦的 Token 序列并应用统一压缩，未能考虑到推理步骤重要性差异显著的层次化结构。

为了克服这一问题，作者引入了思维感知注意力匹配（TAM），这是一种通过三个核心机制利用结构化推理动态的新方法： 1. 思维分割（Thought segmentation），将轨迹分解为不同的推理块。 2. 自适应预算分配（Adaptive budget allocation），根据每个段的大小和相对重要性分配合适的压缩预算。 3. 关键 Token 保护（Pivotal token protection），以保留关键的高注意力推理锚点。

理论证明表明，该分配规则在凸误差模型下保持最优，并且顺序压缩下的累积误差受到严格限制。在 AIME 2024 和 MATH-500 上使用 Qwen3-4B 进行的实证评估表明，在相同的内存占用下，TAM 相比统一压缩提高了任务准确率。此外，周期性 compaction 将峰值内存消耗限制在 3.1–3.2 GB（减少了 65%），同时保持了具有竞争力的准确率。

Paper Metadata / 论文元数据

  • arXiv ID: arXiv:2608.12331 [cs.CL]
  • 主要分类: 计算与语言 (cs.CL)
  • 次要分类: 人工智能 (cs.AI)
  • 提交时间: 2026年6月1日
  • 作者: Yang Liu, Bin Chong, Chongyang Zhang, Hao Zheng, Jiayu Liang, Xu Kefu
  • 许可证: 知识共享署名 4.0 license icon查看许可证

Abstract / 摘要

Reasoning language models generate lengthy chain-of-thought (CoT) sequences whose key-value (KV) cache grows linearly and becomes a memory bottleneck during decoding. Existing compaction methods treat reasoning trajectories as flat token sequences and apply uniform compression, ignoring the hierarchical structure of CoT reasoning where different steps vary drastically in importance. We propose Thought-Aware Attention Matching (TAM), which exploits this structure through three mechanisms:

(i) thought segmentation that decomposes the trajectory into reasoning blocks, (ii) adaptive budget allocation that assigns compression budget based on each segment's importance and size, and (iii) pivotal token protection that preserves high-attention reasoning anchors.

We prove that the allocation rule is optimal under a convex error model and that cumulative error under sequential compaction remains bounded. Experiments on AIME 2024 and MATH-500 with Qwen3-4B show that TAM improves accuracy over uniform compaction at the same memory footprint, with periodic compaction bounding peak memory to 3.1–3.2 GB (a 65% reduction) while maintaining competitive accuracy.

推理语言模型会生成冗长的思维链（CoT）序列，其键值（KV）缓存呈线性增长，在解码过程中成为内存瓶颈。现有的压缩方法将推理轨迹视为平坦的 Token 序列并应用统一压缩，忽略了 CoT 推理的层次化结构，即不同步骤的重要性存在巨大差异。我们提出了思维感知注意力匹配（TAM），它通过以下三种机制利用这种结构：

(i) 思维分割，将轨迹分解为推理块； (ii) 自适应预算分配，根据每个段的重要性和大小分配压缩预算； (iii) 关键 Token 保护，以保留高注意力的推理锚点。

我们证明了该分配规则在凸误差模型下是最优的，并且顺序压缩下的累积误差保持在有界范围内。在 AIME 2024 和 MATH-500 上使用 Qwen3-4B 进行的实验表明，在相同的内存占用下，TAM 提高了相比统一压缩的准确率，通过周期性压缩将峰值内存限制在 3.1–3.2 GB（减少了 65%），同时保持了具有竞争力的准确率。

Key Contributions & Mechanisms / 核心贡献与机制

  • Hierarchical Structural Awareness: Moves beyond flat-sequence token dropping by recognizing the distinct structural phases inherent in chain-of-thought (CoT) reasoning.
  • TAM Framework:
  • Thought Segmentation: Automatically parses reasoning trajectories into modular logical blocks.
  • Adaptive Budget Allocation: Dynamically provisions compaction allowances relative to segment significance.
  • Pivotal Token Protection: Safeguards critical attention anchors essential for complex multi-step deductions.
  • Theoretical Guarantees: Establishes mathematical optimality under a convex error model, ensuring stable error bounds across sequential compaction rounds.
  • Substantial Memory Savings: Achieves a 65% reduction in peak memory usage (down to 3.1–3.2 GB on Qwen3-4B) while outperforming uniform baseline methods in reasoning accuracy on benchmarks like AIME 2024 and MATH-500.
  • 层次化结构感知： 通过识别思维链（CoT）推理中固有的不同结构阶段，超越了扁平序列的 Token 丢弃方式。
  • TAM 框架：
  • 思维分割： 自动将推理轨迹解析为模块化的逻辑块。
  • 自适应预算分配： 根据段落的重要性动态配置压缩配额。
  • 关键 Token 保护： 守护对于复杂多步推演至关重要的关键注意力锚点。
  • 理论保证： 在凸误差模型下建立了数学最优性，确保了连续压缩轮次中的稳定误差边界。
  • 显着的内存节省： 实现了峰值内存使用量 65% 的减少（在 Qwen3-4B 上降至 3.1–3.2 GB），同时在 AIME 2024 和 MATH-500 等基准测试的推理准确率上优于统一的基线方法。

Links & Full-Text Access / 链接与全文访问